'''
Author: SlytherinGe
LastEditTime: 2021-11-03 21:10:10
'''
import os
import cv2
import numpy as np
import math
from fbpca import pca
from area_of_interest.cfar_aoi_mask_generate import CfarBasedAOIMaskGenerator

class PseudoRotetedBoxGenerator(CfarBasedAOIMaskGenerator):

    def __init__(self, anno_root, img_root, cfar_aoi_buff_root, len_ratio):
        super().__init__(anno_root, img_root)
        self.cfar_aoi_buff_root = cfar_aoi_buff_root
        self.cfar_aoi_buff = os.listdir(cfar_aoi_buff_root)
        self.cfar_aoi_buff.sort()
        self.len_ratio = len_ratio
    
    def _generate_rbox(self, index):
        '''
        description: 
        param {*}
        return {numpy array of rbox which shape is [N, 5], (cx, cy, w, h, theta)}
        '''
        index = index
        anno = self.get_anno_from_index(index)
        bboxes = []
        for bndbox in anno['object']:
            bbox = (int(bndbox['bndbox']['xmin']), int(bndbox['bndbox']['ymin']), int(bndbox['bndbox']['xmax']), int(bndbox['bndbox']['ymax']))
            bboxes.append(bbox)
        bboxes = np.array(bboxes)
        shape = (int(anno['size']['height']), int(anno['size']['width']), int(anno['size']['depth']))
        results = dict(gt_bboxes=bboxes, img_shape=shape)
        img_shape = results['img_shape']
        gt_bboxes = results['gt_bboxes'].astype(np.intp)
        if self.cfar_aoi_buff_root is not None:
            cfar_buffer_file = os.path.join(self.cfar_aoi_buff_root, self.cfar_aoi_buff[index])
            cfar_aoi = cv2.imread(cfar_buffer_file)[:,:,0]
        else:
            # TODO:runtime cfar
            img = cv2.imread(os.path.join(self.img_root, self.img_files[index]))
        # start using pca to generate
        # results is a list of tuples: (center(x,y), principle component)
        results = []

        single_ship_buff = np.zeros((img_shape[0], img_shape[1]))
        for i in range(gt_bboxes.shape[0]):
            single_ship_buff[gt_bboxes[i,1]:gt_bboxes[i,3]+1, gt_bboxes[i,0]:gt_bboxes[i,2]+1] = cfar_aoi[gt_bboxes[i,1]:gt_bboxes[i,3]+1, gt_bboxes[i,0]:gt_bboxes[i,2]+1]
            points = np.nonzero(single_ship_buff>128)
            # (x, y)
            points_stacked = np.vstack((points[1],points[0])).T
            # PCA
            U, s, Va = pca(points_stacked, 2)
            # theta
            theta = math.atan2(Va[0,1], Va[0,0])
            # get target center
            center = np.floor((gt_bboxes[i,2:] + gt_bboxes[i,:2])/2+1).astype(np.int)
            # cal w, h
            x_dim = np.sort(U[:,0]*s[0])
            y_dim = np.sort(U[:,1]*s[1])
            ylen = len(y_dim)
            xlen = len(x_dim)
            ystart = int(self.len_ratio*ylen)
            yend = 1 if (ystart==0) else ystart 
            xstart = int(self.len_ratio*xlen)
            xend = 1 if (xstart==0) else xstart 
            h = y_dim[-yend] - y_dim[ystart]
            w = x_dim[-xend] - x_dim[xstart]
            single_ship_buff = np.zeros((img_shape[0], img_shape[1]))
            # TODO: modify class
            results.append(np.array([1,1.0,center[0],center[1],w, h, theta*180/math.pi]))

        results = np.vstack(results)
        
        return results


